#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""            
@FileName: dft_interp.py
@Version: 1.0
@Date: 2023/7/31 11:09
@Description:
"""

import numpy as np
from typing import Tuple
from matplotlib import pyplot as plt


# 程序 10.1 傅立叶插值
# 使用三角函数 P(t) 插值区间 [c,d] 上的 n 个点
# 并画出在 p(>=n) 个均匀分布的点上的插值
# 输入: 区间 [c,d], 数据点 x, 数据个数n, 偶数 p>=n
# 输出: 插值得到的数据点 xp
def dft_interp(inter: Tuple[float, ...], x: np.ndarray, n: int, p: int):
    c = inter[0]
    d = inter[1]
    t = c + (d - c) * np.arange(0, n) / n
    tp = c + (d - c) * np.arange(0, p) / p
    # 应用 DFT
    y = np.fft.fft(x)
    # yp 保存 ifft 的系数
    yp = np.zeros(p, dtype=complex)
    # 将频率由 n 改成 p
    yp[0:n // 2 + 1] = y[0:n // 2 + 1]
    #  对上半部分做同样处理
    yp[p - n // 2 + 1:p] = y[n // 2 + 1:n]
    # 求逆，重建数据
    xp = np.real(np.fft.ifft(yp) * (float(p / n)))
    # 画出数据点以及对应的插值
    plt.plot(t, x, 'o', tp, xp)
    plt.show()


def test_np_fft():
    x = np.array([0, 1, 2, 3, 4])
    a = np.fft.fft(x, norm='backward')
    a_real = np.fft.rfft(x, norm='backward')
    freq_lst = np.fft.fftfreq(len(x))
    a_shift = np.fft.fftshift(a)

    print(f'orig: {x}')
    print(f'tranformed complex: {a}')
    print(f'tranformed real: {a_real}')
    print(f'freq list: {freq_lst}')
    print(f'trans shifted: {a_shift}')


if __name__ == '__main__':
    irng = (0, 1)
    xlst = np.array([-2.2, -2.8, -6.1, -3.9, 0.0, 1.1, -0.6, -1.1])
    dft_interp(irng, xlst, 8, 100)
